package com.xwj.nio.tcp;

import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.net.InetSocketAddress;
import java.net.ServerSocket;
import java.nio.ByteBuffer;
import java.nio.channels.*;
import java.nio.charset.Charset;
import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;

/**
 * nio通道channel之ServerSocketChannel、Selector
 * TCP协议
 *
 * @author xwj
 * @date 2020/6/10
 */
public class NioReceiveServer {
    private static Map<SelectableChannel, Client> clientMap = new HashMap<>(64);

    private static ByteBuffer buf = ByteBuffer.allocate(2048);

    public static void main(String[] args) throws IOException {
        NioReceiveServer server = new NioReceiveServer();
        server.startServer();
    }

    private static class Client {
        String fileName;
        Long fileLen;
        InetSocketAddress remoteAddress;
        FileChannel fileChannel;
    }

    public void startServer() throws IOException {
        Selector selector = Selector.open();
        ServerSocketChannel ssc = ServerSocketChannel.open();
        ssc.configureBlocking(false);
        ssc.bind(new InetSocketAddress(8888));
        ssc.register(selector, SelectionKey.OP_ACCEPT);
        System.out.println("ServerSocketChannel 监听中......");
        while (selector.select() > 0) {
            Iterator<SelectionKey> selectionKey = selector.selectedKeys().iterator();
            while (selectionKey.hasNext()) {
                SelectionKey key = selectionKey.next();
                if (key.isAcceptable()) {
                    // 新连接事件，获取客户端连接
                    ServerSocketChannel server = (ServerSocketChannel) key.channel();
                    SocketChannel socketChannel = server.accept();
                    if (socketChannel == null) {
                        continue;
                    }
                    // 客户端新连接切换为非阻塞模式
                    socketChannel.configureBlocking(false);
                    socketChannel.register(selector, SelectionKey.OP_READ);
                    Client client = new Client();
                    client.remoteAddress = (InetSocketAddress) socketChannel.getRemoteAddress();
                    clientMap.put(socketChannel, client);
                    System.out.println(socketChannel.getRemoteAddress() + "连接成功");
                } else if(key.isReadable()) {
                    processData(key);
                }
                // 防止下一次被Select函数选中
                selectionKey.remove();
            }
        }
    }

    /**
     * 处理客户端传输的数据
     * @param key
     */
    private void processData(SelectionKey key) {
        Client client = clientMap.get(key.channel());
        SocketChannel socketChannel = (SocketChannel) key.channel();
        int len = 0;
        try {
            buf.clear();
            while((len = socketChannel.read(buf)) != -1) {
                // 切换成可读模式
                buf.flip();
                if(client.fileName == null) {
                    // 接收文件名
                    String fileName = Charset.forName("UTF-8").decode(buf).toString();
                    File target = new File("C:\\Users\\Administrator\\Desktop\\target\\" + fileName);
                    FileChannel fileChannel = new FileOutputStream(target).getChannel();
                    client.fileChannel = fileChannel;
                    client.fileName = fileName;
                } else if(client.fileLen == null) {
                    long fileLen = buf.getLong();
                    client.fileLen = fileLen;
                } else {
                    client.fileChannel.write(buf);
                }
                buf.clear();
            }
            key.cancel();
        } catch (Exception e) {
            key.cancel();
            e.printStackTrace();
            return;
        }
        if(len == -1) {
            try {
                client.fileChannel.close();
                key.cancel();
                System.out.println("文件接收成功：" + client.fileName);
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
    }
}
